package com.github.prontera.common.auth0.jwt;

import java.nio.charset.StandardCharsets;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

import org.apache.commons.codec.binary.Base64;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.MapperFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.github.prontera.common.auth0.jwt.algorithms.Algorithm;
import com.github.prontera.common.auth0.jwt.exceptions.JWTCreationException;
import com.github.prontera.common.auth0.jwt.exceptions.SignatureGenerationException;
import com.github.prontera.common.auth0.jwt.impl.ClaimsHolder;
import com.github.prontera.common.auth0.jwt.impl.PayloadSerializer;
import com.github.prontera.common.auth0.jwt.impl.PublicClaims;
import com.github.prontera.util.DateTimeUtil;

public final class JWTCreator {

	private final Algorithm algorithm;
	private final String headerJson;
	private final String payloadJson;

	private JWTCreator(Algorithm algorithm, Map<String, Object> headerClaims, Map<String, Object> payloadClaims) throws JWTCreationException {
		this.algorithm = algorithm;
		try {
			ObjectMapper mapper = new ObjectMapper();
			SimpleModule module = new SimpleModule();
			module.addSerializer(ClaimsHolder.class, new PayloadSerializer());
			mapper.registerModule(module);
			mapper.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true);
			headerJson = mapper.writeValueAsString(headerClaims);
			payloadJson = mapper.writeValueAsString(new ClaimsHolder(payloadClaims));
		} catch (JsonProcessingException e) {
			throw new JWTCreationException("Some of the Claims couldn't be converted to a valid JSON format.", e);
		}
	}

	static JWTCreator.Builder init() {
		return new Builder();
	}

	public static class Builder {
		private final Map<String, Object> payloadClaims;
		private Map<String, Object> headerClaims;

		Builder() {
			this.payloadClaims = new HashMap<String, Object>();
			this.headerClaims = new HashMap<String, Object>();
		}

		/**
		 * 添加自定义header
		 */
		public Builder withHeader(Map<String, Object> headerClaims) {
			this.headerClaims = new HashMap<String, Object>(headerClaims);
			return this;
		}

		/**
		 * Add a specific Key Id ("kid") claim to the Header.
		 *
		 * @param keyId
		 *            the Key Id value.
		 * @return this same Builder instance.
		 */
		public Builder withKeyId(String keyId) {
			this.headerClaims.put(PublicClaims.KEY_ID, keyId);
			return this;
		}

		/**
		 * 添加发行人
		 */
		public Builder withIssuer(String issuer) {
			addClaim(PublicClaims.ISSUER, issuer);
			return this;
		}

		/**
		 * Add a specific Subject ("sub") claim to the Payload.
		 *
		 * @param subject
		 *            the Subject value.
		 * @return this same Builder instance.
		 */
		public Builder withSubject(String subject) {
			addClaim(PublicClaims.SUBJECT, subject);
			return this;
		}

		/**
		 * Add a specific Audience ("aud") claim to the Payload.
		 *
		 * @param audience
		 *            the Audience value.
		 * @return this same Builder instance.
		 */
		public Builder withAudience(String... audience) {
			addClaim(PublicClaims.AUDIENCE, audience);
			return this;
		}

		/**
		 * Add a specific Expires At ("exp") claim to the Payload.
		 *
		 * @param expiresAt
		 *            the Expires At value.
		 * @return this same Builder instance.
		 */
		public Builder withExpiresAt(Date expiresAt) {
			addClaim(PublicClaims.EXPIRES_AT, expiresAt);
			return this;
		}

		/**
		 * Add a specific Expires At ("exp") claim to the Payload.
		 *
		 * @param expiresAt
		 *            the Expires At value.
		 * @return this same Builder instance.
		 */
		public Builder withExpiresAt(Integer expiresAtSeconds) {
			addClaim(PublicClaims.EXPIRES_AT, DateTimeUtil.addSeconds(DateTimeUtil.getSystemDate(), expiresAtSeconds));
			return this;
		}

		/**
		 * Add a specific Not Before ("nbf") claim to the Payload.
		 *
		 * @param notBefore
		 *            the Not Before value.
		 * @return this same Builder instance.
		 */
		public Builder withNotBefore(Date notBefore) {
			addClaim(PublicClaims.NOT_BEFORE, notBefore);
			return this;
		}

		/**
		 * Add a specific Issued At ("iat") claim to the Payload.
		 *
		 * @param issuedAt
		 *            the Issued At value.
		 * @return this same Builder instance.
		 */
		public Builder withIssuedAt(Date issuedAt) {
			addClaim(PublicClaims.ISSUED_AT, issuedAt);
			return this;
		}

		/**
		 * Add a specific JWT Id ("jti") claim to the Payload.
		 *
		 * @param jwtId
		 *            the Token Id value.
		 * @return this same Builder instance.
		 */
		public Builder withJWTId(String jwtId) {
			addClaim(PublicClaims.JWT_ID, jwtId);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(String name, Boolean value) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, value);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(String name, Integer value) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, value);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(String name, Double value) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, value);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(Map<String, Object> claims) throws IllegalArgumentException {
			for (String claim : claims.keySet()) {
				assertNonNull(claim);
				addClaim(claim, claims.get(claim));
			}
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(String name, String value) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, value);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withClaim(String name, Date value) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, value);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withArrayClaim(String name, String[] items) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, items);
			return this;
		}

		/**
		 * 添加自定义值
		 */
		public Builder withArrayClaim(String name, Integer[] items) throws IllegalArgumentException {
			assertNonNull(name);
			addClaim(name, items);
			return this;
		}

		/**
		 * 签名
		 */
		public String sign(Algorithm algorithm) throws IllegalArgumentException, JWTCreationException {
			if (algorithm == null) {
				throw new IllegalArgumentException("The Algorithm cannot be null.");
			}
			headerClaims.put(PublicClaims.ALGORITHM, algorithm.getName());
			headerClaims.put(PublicClaims.TYPE, "JWT");
			return new JWTCreator(algorithm, headerClaims, payloadClaims).sign();
		}

		private void assertNonNull(String name) {
			if (name == null) {
				throw new IllegalArgumentException("The Custom Claim's name can't be null.");
			}
		}

		private void addClaim(String name, Object value) {
			if (value == null) {
				payloadClaims.remove(name);
				return;
			}
			payloadClaims.put(name, value);
		}
	}

	private String sign() throws SignatureGenerationException {
		String header = Base64.encodeBase64URLSafeString((headerJson.getBytes(StandardCharsets.UTF_8)));
		String payload = Base64.encodeBase64URLSafeString((payloadJson.getBytes(StandardCharsets.UTF_8)));
		String content = String.format("%s.%s", header, payload);
		byte[] signatureBytes = algorithm.sign(content.getBytes(StandardCharsets.UTF_8));
		String signature = Base64.encodeBase64URLSafeString((signatureBytes));
		return String.format("%s.%s", content, signature);
	}
}
